"""
数据库连接和会话管理模块
"""
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, Session
from sqlalchemy.pool import StaticPool
from sqlalchemy import text
from contextlib import contextmanager
from typing import Generator
import logging

from .config import settings
from models.base import Base
from models import *

logger = logging.getLogger(__name__)


# 创建数据库引擎
def create_database_engine(database_url: str = None):
    """创建数据库引擎"""
    if database_url is None:
        database_url = settings.database_url

    # SQLite特定配置
    if database_url.startswith("sqlite"):
        engine = create_engine(
            database_url,
            connect_args={
                "check_same_thread": False,
                "timeout": 30,  # 增加超时时间
                "isolation_level": None,  # 自动提交模式
            },
            poolclass=StaticPool,
            echo=settings.debug,
            pool_pre_ping=True,  # 连接前检查
            pool_recycle=3600,  # 1小时回收连接
        )
    else:
        # PostgreSQL/MySQL配置
        engine = create_engine(
            database_url,
            echo=settings.debug,
            pool_size=20,
            max_overflow=30,
            pool_pre_ping=True,
            pool_recycle=3600,
        )

    return engine


# 创建数据库引擎实例
engine = create_database_engine()


# 创建会话工厂
SessionLocal = sessionmaker(
    autocommit=False,
    autoflush=False,
    bind=engine
)


def get_db() -> Generator[Session, None, None]:
    """获取数据库会话"""
    db = SessionLocal()
    try:
        yield db
        db.commit()  # 正常提交
    except Exception as e:
        logger.error(f"数据库会话错误: {e}")
        db.rollback()
        raise
    finally:
        db.close()


@contextmanager
def get_db_session():
    """获取数据库会话上下文管理器"""
    db = SessionLocal()
    try:
        yield db
        db.commit()
    except Exception as e:
        logger.error(f"数据库会话错误: {e}")
        db.rollback()
        raise
    finally:
        db.close()


def create_tables():
    """创建数据库表"""
    try:
        from models.user import User
        from models.service_class import ServiceClass, ServiceInstance, ServiceUsageLog, ServiceConfigTemplate
        from models.prompt_template import PromptTemplate, PromptParameter, PromptUsageRecord
        from models.analysis_task import AnalysisTask, TaskService, TaskTemplate
        from models.report import Report, ReportTemplate, ReportSchedule
        from models.system import SystemConfig, OperationLog, SystemHealth, Notification, BackupRecord
        from models.petition_record import (PetitionRecord, RegionalAnalysisResult,
                                            ComplaintTypeAnalysisResult, SentimentAnalysisResult,
                                            ComplianceAnalysisResult, DuplicateAnalysisResult)
        Base.metadata.create_all(bind=engine)
        logger.info("数据库表创建成功")
    except Exception as e:
        logger.error(f"创建数据库表失败: {e}")
        raise


def drop_tables():
    """删除数据库表"""
    try:
        Base.metadata.drop_all(bind=engine)
        logger.info("数据库表删除成功")
    except Exception as e:
        logger.error(f"删除数据库表失败: {e}")
        raise


def check_database_connection() -> bool:
    """检查数据库连接"""
    try:
        with get_db_session() as db:
            # 执行简单查询
            db.execute(text("SELECT 1"))
        logger.info("数据库连接正常")
        return True
    except Exception as e:
        logger.error(f"数据库连接失败: {e}")
        return False


# 数据库会话依赖
class DatabaseManager:
    """数据库管理器"""
    
    def __init__(self):
        self.engine = engine
        self.session_factory = SessionLocal
    
    def get_session(self) -> Session:
        """获取新的数据库会话"""
        return self.session_factory()
    
    def close_all_connections(self):
        """关闭所有数据库连接"""
        self.engine.dispose()
    
    def execute_raw_sql(self, sql: str, params: dict = None):
        """执行原始SQL"""
        with get_db_session() as db:
            result = db.execute(sql, params or {})
            return result.fetchall()


# 全局数据库管理器实例
db_manager = DatabaseManager()


# 便捷函数
def db_session():
    """数据库会话装饰器"""
    def decorator(func):
        def wrapper(*args, **kwargs):
            with get_db_session() as db:
                return func(db, *args, **kwargs)
        return wrapper
    return decorator